# @Time : 2021/9/28 10:55
# @Author : xiashuobad
# @File : image_utils.py
# @Software: PyCharm
########################
#  _ooOoo_
# o8888888o
# 88" . "88
# (| -_- |)
# O\ = /O
# ___/`---'\____
# .   ' \\| |// `.
# / \\||| : |||// \
# / _||||| -:- |||||- \
# | | \\\ - /// | |
# | \_| ''\---/'' | |
# \ .-\__ `-` ___/-. /
# ___`. .' /--.--\ `. . __
# ."" '< `.___\_<|>_/___.' >'"".
# | | : `- \`.;`\ _ /`;.`/ - ` : | |
# \ \ `-. \_ __\ /__ _/ .-` / /
# ======`-.____`-.___\_____/___.-`____.-'======
# '=---=' bug泛滥 佛已瘫痪
import contextlib
import json
import re
# import mmcv
# import warnings
import numpy as np
# from rasterio.features import sieve
from rasterio.mask import mask
from PIL import Image
from osgeo import ogr, osr
# from skimage import morphology
from tqdm import tqdm
from functools import reduce
from rasterio import features
from constant import DICT_FIELD_TYPE
from file_utils import MyPath
# from parallel import parallel_run_func
from open_image import OpenImage, cv2
# from tools.postprocess.rdp import rdp
from transform_types import adjust_transform, geo_transform2affine, \
    format_colormap2pil, format_class_info


def read_image_as_array(img_path: str, backend='gdal', **kwargs) -> np.ndarray:
    """
    gdal读取图像，返回图像array以及坐标和投影
    :param img_path: 图片路径
    :param backend：使用的后端图片库，支持cv2，pil，gdal，rasterio
    :return:
    """
    with OpenImage(img_path, backend=backend, **kwargs) as img:
        img_array = img.img_array
    return img_array


# def save_mask2color(_mask: (np.ndarray, str), save_path, color_map: (list, dict), dict_replace_values=None,
#                     backend='pil', geo_transform=None, project=None):
#     """
#     将标签图像array或者预测图像array转换成伪彩色图像并保存
#     :param _mask: mask图像路径或者array
#     :param save_path: 伪彩色标签保存路径
#     :param color_map: rgb颜色映射表,支持pil，rasterio以及二维列表三个格式
#     :param dict_replace_values: 字典类型！key为原始值(多个值可以用元组)，value为新值(长度为2的元组，第一个元素为key值对应的新值，
#     第二个值表示除key值外的其他值对应的新值,默认为None表示不变)'
#     :param backend: 使用的后端图片库，'pil','gdal','rasterio'
#     :param geo_transform:坐标系数
#     :param project:坐标系
#     :return:
#     """
#     if dict_replace_values is None:
#         dict_replace_values = {}
#     if isinstance(_mask, str):  # 如果mask参数传的是文件路径，则先读取图像
#         _mask = read_image_as_array(_mask, backend=backend)
#     if len(_mask.shape) == 3:  # 非单通道灰度图，取第一个通道转为灰度图
#         _mask = reshape_as_raster(_mask)[0, ...]
#     assert len(_mask.shape) == 2, '只支持单通道灰度图！'
#     _mask = img_replace_values(_mask, dict_replace_values)
#     _mask = _mask.astype(np.uint8)  # 转uint8图片格式
#     write_image(_mask, save_path, geo_transform, project, color_map=color_map, backend=backend)
#
#
# def batch_mask2color(mask_dir, color_label_dir, color_map, dict_replace_values=None, num_process=4):
#     """
#     批量将mask_dir目录的所有原始标签转换成伪彩色标签并保存
#     :param mask_dir:
#     :param color_label_dir:
#     :param color_map:
#     :param dict_replace_values: 字典类型！key为原始值(多个值可以用元组)，value为新值(长度为2的元组，第一个元素为key值对应的新值，第二个值表示除key值外的其他值对应的新值,默认为None表示不变)'
#     :param num_process: 使用进程的数量，为空或者其他假值，表示根据机器核心数自动设置
#     :return:
#     """
#     MyPath(color_label_dir).makedirs(exist_ok=True)
#     iterable_args = [(mask_path.path, mask_path.replace_parent(color_label_dir).path) for mask_path in
#                      MyPath(mask_dir).iter_dir(ImageType.pattern())]
#     parallel_run_func(partial(save_mask2color, color_map=color_map, dict_replace_values=dict_replace_values),
#                       iterable_args, num_workers=num_process, desc='批量转换伪彩色标签图片中。。。')


def image_slide_cutting(end_x, end_y, start_x=0, start_y=0, window_size=512, overlapping_pixels=0):
    """
    大影像滑动窗口切割生成器，返回当前切割得图片左上角得像素坐标(x,y)和实际切割的图像的尺寸(window_size_x, window_size_y)
    :param end_x: x上界
    :param end_y: y上界
    :param start_x: x下界
    :param start_y: y下界
    :param window_size: 切割的图片尺寸大小,(宽，高)，为整数时，宽高相同
    :param overlapping_pixels: 窗口滑动时，重叠的像素大小。
    :return:
    """
    assert end_x > start_x and end_y > start_y
    if isinstance(window_size, int):
        window_size_x = window_size_y = window_size
    else:
        assert isinstance(window_size, (list, tuple))
        window_size_x, window_size_y = window_size[:2]
    # 存在切割尺寸大于原图尺寸的情况，特殊处理
    window_size_y = min(window_size_y, end_y - start_y)
    window_size_x = min(window_size_x, end_x - start_x)
    stride_x = window_size_x - overlapping_pixels
    stride_y = window_size_y - overlapping_pixels
    # 按照从左往右，从上往下，逐行扫描
    for y in list(range(start_y, end_y - window_size_y, stride_y)) + [end_y - window_size_y]:
        for x in list(range(start_x, end_x - window_size_x, stride_x)) + [end_x - window_size_x]:
            yield (x, y), (window_size_x, window_size_y)


# def slice_generator(image: (str, OpenImage), window_size: (int, tuple, list), overlapping_pixels=0,
#                     start_coordinate: (tuple, list) = (0, 0), end_coordinate: (tuple, list) = None) -> (
#         np.ndarray, (int, int), (int, int)):
#     """
#     图像滑动窗口切割生成器，返回当前切割得图片左上角得像素坐标(x,y)和实际切割的图像的尺寸(window_size_x, window_size_y)
#     :param image: 图片array
#     :param window_size: 切割的图片尺寸大小,(宽，高)，为整数时，宽高相同
#     :param overlapping_pixels: 窗口滑动时，重叠的像素大小。
#     :param start_coordinate:
#     :param end_coordinate:
#     :return: 返回切片的array,(x,y),(window_size_x, window_size_y)
#     """
#     if isinstance(image, str):
#         image = OpenImage(image)
#     assert isinstance(image, OpenImage), 'image参数应为图片路径或者OpenImage对象！'
#     end_coordinate = end_coordinate or (image.width, image.height)
#
#     for (x, y), (window_size_x, window_size_y) in tqdm(
#             list(image_slide_cutting(*end_coordinate, *start_coordinate, window_size,
#                                      overlapping_pixels))):
#         yield image.read_window(x, y, window_size_x, window_size_y), (x, y), (window_size_x, window_size_y)
#     image.close()


def create_field_defn(out_layer, class_fields: dict):
    """
    给据所给字段属性列表给图层创建属性字段
    Args:
        out_layer: 图层
        class_fields: 类别字段，{字段名称:字段类型}

    Returns:

    """
    for field, field_type in class_fields.items():
        fd = ogr.FieldDefn(field, DICT_FIELD_TYPE.get(field_type, ogr.OFTString))
        out_layer.CreateField(fd)


def shp2raster(raster_path: str, shp: (str, list), save_tif_path, field_raster: str, field_mapping=None,
               color_mapping=None, attr_filter: str = None):
    """
    shp 转 栅格图像
    :param raster_path: tif 原图，用来提供坐标信息
    :param shp: 待转的 shp文件路径或者postgis数据库连接地址,或者直接shapes对象:[(geometry,id),...]
    :param save_tif_path: 转好的 tif 保存路径
    :param field_raster: 用来匹配的字段名称
    :param field_mapping: 用来匹配信息的字段映射，dict
    :param color_mapping: rgb颜色映射表，rasterio格式或pil格式
    :param attr_filter: 属性过滤条件查询语句
    """
    if field_mapping is None:
        field_mapping = {}
    with OpenImage(raster_path) as img_ds:
        geo_transform = img_ds.geo_transform
        projection = img_ds.project
        width, height = img_ds.size
    geo_transform = geo_transform2affine(geo_transform)  # 转换为rasterio格式
    if isinstance(shp, str):
        with read_shp(shp, attr_filter) as layer:
            shp = [(json.loads(feature.geometry().ExportToJson()),
                    field_mapping[feature.GetField(field_raster)] if field_mapping else feature.GetField(field_raster))
                   for feature in layer]

    mask_array = features.rasterize(shp, (height, width), transform=geo_transform, all_touched=True)
    write_image(mask_array, save_tif_path, geo_transform, projection, color_map=color_mapping)


@contextlib.contextmanager
def read_shp(shp_path: str, attr_filter=None, spatial_filter=None) -> ogr.Layer:
    """
    读取shp，返回layer，支持shpfile文件路径和postgis数据库读取
    :param shp_path: 数据库连接地址格式为 "PG:host=localhost port=5432 dbname=kqgis user=postgres password=135456
    table=dawenhe-clip name='水域' or name='房屋建筑'"
    数据库连接地址格式分为3个部分："PG:host=localhost port=5432 dbname=kqgis user=postgres password=135456","table=...",
    "name=...",后面两个分别表示查询的表名和属性过滤条件
    :param attr_filter: 属性过滤条件表达式，字符串，例如："name='水域' or name='房屋建筑'"
    :param spatial_filter: 空间过滤，Geometry类型，给定空间范围
    :return:
    """
    assert isinstance(shp_path, str)
    if MyPath(shp_path).is_file():  # shp文件路径
        ds: ogr.DataSource = ogr.Open(shp_path)
        layer = ds.GetLayer()
    elif shp_path.startswith('PG:') and 'table' in shp_path:  # postgis数据库地址
        res = re.findall(r'\s*(PG:.*)\s+table=(\S+)\s*', shp_path.strip())
        assert len(res) > 0, f'{shp_path}格式错误！'
        db_path, table = res[0]
        ds: ogr.DataSource = ogr.Open(db_path)
        layer: ogr.Layer = ds.GetLayerByName(table)
    else:
        raise Exception(f'{shp_path} 格式错误！')
    attr_filter and layer.SetAttributeFilter(attr_filter)
    spatial_filter and layer.SetSpatialFilter(spatial_filter)
    assert len(layer) > 0, f'{shp_path} 读取错误！'
    yield layer
    del ds


def write_image(img: (np.ndarray, str), save_path, geo_transform=None, projection_ref=None, ref_raster_path=None,
                color_map: dict = None, nodata=None, backend='gdal'):
    """
    图像array保存为带坐标的tif栅格数据
    :param img:
    :param save_path: 保存路径
    :param geo_transform: 地理坐标变换系数
    :param projection_ref: 坐标系
    :param ref_raster_path: 参考栅格影像路径
    :param color_map: rgb颜色映射表，字典类型
    :param nodata: 无效值
    :param backend: 'gdal','rasterio','pil','cv2'
    :return:
    """
    if isinstance(img, str):  # 传入的是图片路径时，先读取图片
        with OpenImage(img, backend=backend) as ds_img:
            img = ds_img.img_array
            geo_transform = geo_transform or ds_img.geo_transform  # 当参数设置了geo_transform的值就以参数，否则以读取图片里的值
            projection_ref = projection_ref or ds_img.project  # 同上
            nodata = nodata or ds_img.nodata  # 同上
    if ref_raster_path and MyPath(ref_raster_path).exists():  # 当设置了参考栅格的路径时优先级最高，使用参考栅格的坐标变换系数和坐标系
        with OpenImage(ref_raster_path, backend=backend) as ds_img:
            geo_transform = ds_img.geo_transform
            projection_ref = ds_img.project
            nodata = ds_img.nodata
    MyPath(save_path).make_parent_dirs()  # 提前创建父目录，防止写入文件报错
    with OpenImage(save_path, 'w', backend) as ds_img:  # 写入数据到图片文件
        ds_img.set_img_array(img).set_nodata(nodata).set_geo_transform(geo_transform).set_project(
            projection_ref).set_color_map(color_map)


def mask2shp(_mask: (np.ndarray, str), save_shp_path: str, class_info: (list, dict), geo_transform=None,
             projection_ref=None, ref_raster_path=None, del_background=False, simplify_tolerance=2):
    """
    mask预测图array转矢量shp
    :param _mask: mask图片路径或者array
    :param save_shp_path: 矢量shp文件保存路径
    :param class_info: 分类信息，格式同dict_model配置文件
    :param geo_transform: 地理坐标变换系数
    :param projection_ref: 投影信息
    :param ref_raster_path: 参考影像
    :param del_background: 是否删除背景0
    :param simplify_tolerance: 边缘点精简，容忍度参数，值越大精简越狠
    :return: 返回所有feature（目标区域）的列表
    """
    if not isinstance(class_info, dict):
        class_info = format_class_info(class_info)
    if isinstance(_mask, str):  # mask为文件路径时，先读取mask图片文件
        with OpenImage(_mask) as img:
            _mask = img.img_array
            geo_transform = geo_transform or img.geo_transform  # 参数未设置地理坐标仿射变换和坐标系时，直接从mask图像文件中获取
            projection_ref = projection_ref or img.project
    assert len(_mask.shape) == 2, 'mask_array应为单通道灰度图！'
    if ref_raster_path and MyPath(ref_raster_path).exists():  # 当设置了参考栅格影像的路径的时，地理坐标和坐标系以参考栅格影像的为准
        with OpenImage(ref_raster_path) as img:
            geo_transform = img.geo_transform
            projection_ref = img.project
    MyPath(save_shp_path).make_parent_dirs()  # 自动创建结果矢量的目录
    geo_transform = adjust_transform(geo_transform)  # 无坐标系默认变换系数的情况下会上下翻转，自动调整
    geo_transform = geo_transform2affine(geo_transform)  # 转换为rasterio支持的格式affine
    shapes = features.shapes(_mask, mask=(_mask != 0) if del_background else None,
                             transform=geo_transform)  # 栅格矢量化
    drv: ogr.Driver = ogr.GetDriverByName('ESRI Shapefile')  # 创建矢量驱动
    if MyPath(save_shp_path).exists():  # 删除之前的结果
        drv.DeleteDataSource(save_shp_path)  # 这个函数会将.shp文件以及它的依赖文件一起删除
    dst_shp_ds = drv.CreateDataSource(save_shp_path)  # 创建文件资源
    proj = osr.SpatialReference()  # 创建坐标系类
    proj.ImportFromWkt(projection_ref if projection_ref else '')  # 导入坐标系
    dst_layer: ogr.Layer = dst_shp_ds.CreateLayer('layer', geom_type=ogr.wkbPolygon, srs=proj)  # 创建多边形图层，并写入坐标系
    create_field_defn(dst_layer, class_fields := class_info['class_fields'])  # 创建矢量属性字段
    for polygon, class_id in shapes:  # 遍历shapes，每个shape中保存的是要素的类型以及坐标（字典格式），还有id值
        class_id = int(class_id)
        geometry: ogr.Geometry = ogr.CreateGeometryFromJson(json.dumps(polygon))  # 先将字典转为json格式，再从geojson格式创建geometry
        # 简化边缘点
        geometry = geometry.SimplifyPreserveTopology(simplify_tolerance * geo_transform.a)
        feature: ogr.Feature = ogr.Feature(dst_layer.GetLayerDefn())  # 创建feature要素
        for field, field_type in class_fields.items():  # 遍历所有属性字段，逐个设置值
            # 当类别列表中未设置该属性字段的值时，根据矢量字段类型设置默认值
            field_value = class_info['class_names'][class_id].get(field, field_type())
            feature.SetField(field, field_value)  # 设置字段的值
        feature.SetGeometry(geometry)  # 给feature要素添加多边形的坐标点位置信息
        dst_layer.CreateFeature(feature)  # 将feature要素加到图层中

    dst_shp_ds.Destroy()  # 释放资源，更新写入shp文件


def stitching_image(list_img_paths, save_path=None, unit_shape=(512, 512, 3), gap=10):
    """
    将多张图片水平拼接成一张图片
    :param list_img_paths: 待拼接的图片路劲列表
    :param save_path: 拼接后的图片保存路径
    :param unit_shape: 单元格图像的shape大小，（h,w,c）
    :param gap: 间隔大小，像素
    :return: 返回拼接后的图片array
    """
    list_img_arrays = []
    gap_shape = list(unit_shape)
    gap_shape[1] = gap
    gap_array = np.full(tuple(gap_shape), 255)
    for img_path in list_img_paths:
        list_img_arrays.append(gap_array)
        image_array = read_image_as_array(img_path, backend='cv2')
        if image_array.shape != unit_shape:
            image_array = cv2.resize(image_array, unit_shape[:2])
        list_img_arrays.append(image_array)
    list_img_arrays.append(gap_array)
    img_stitched = np.hstack(list_img_arrays)  # 将所有图片进行水平拼接
    save_path and write_image(img_stitched, save_path, backend='cv2')  # 保存路径不为空，则保存到文件
    return img_stitched


# def batch_stitching_image(list_img_dirs, dst_dir, gap=10, unit_shape=(512, 512, 3), num_process=4):
#     """
#     批量对多组（大于等于2）目录下得同名图片进行拼接
#     :param list_img_dirs: 要拼接得图片目录列表
#     :param dst_dir: 拼接后图片输出的目录
#     :param gap: 拼接间隔，单位是像素
#     :param unit_shape: 每个单元格图像的shape，格式为（h,w,c）
#     :param num_process: 多进程个数
#     :return:
#     """
#
#     assert len(list_img_dirs) >= 2, '至少需要2组图片进行拼接！'
#     list_img_suffix = get_list_suffix_from_dirs(list_img_dirs)  # 自动获取各个目录下图片的后缀名
#     MyPath(dst_dir).makedirs(exist_ok=True)
#     iterable_args = []
#     for image_path in MyPath(list_img_dirs[0]).iter_dir(ImageType.pattern()):
#         image_path.real_suffix = list_img_suffix[0]
#         list_img_paths = [image_path.replace_parent(img_dir).replace_suffix(list_img_suffix[i]).path for i, img_dir in
#                           enumerate(list_img_dirs)]
#         save_path = image_path.replace_parent(dst_dir).replace_suffix('.png').path
#         iterable_args.append((list_img_paths, save_path))
#     parallel_run_func(stitching_image, iterable_args, num_workers=num_process, desc='批量拼接图片中。。。')


def write_coordinate_to_raster(target_image_path, geo_transform=None, project=None, reference_image_path=None):
    """
    给目标影像写入地理坐标
    :param reference_image_path:参考图片路径
    :param target_image_path:目标图片路径
    :param geo_transform:坐标变换系数
    :param project:坐标系
    :return:
    """
    if reference_image_path is not None:
        with OpenImage(reference_image_path) as reference_data:
            geo_transform = reference_data.geo_transform
            project = reference_data.project
    with OpenImage(target_image_path, 'u') as target_data:
        target_data.set_geo_transform(geo_transform).set_project(project)


def show_mask_to_rgb(image: (np.ndarray, str), color_map=None):
    """
    简单显示一下图片，单通道灰度图可以设置rgb映射
    Args:
        image: 图片路径或者array
        color_map: rgb映射列表

    Returns:

    """
    if isinstance(image, str):
        image = read_image_as_array(image)
    image = Image.fromarray(image)
    if color_map is not None:
        color_map = format_colormap2pil(color_map)
        image.putpalette(color_map)
    image.show()


# def generate_graph(im):
#     if len(im.shape) == 3:
#         warnings.warn(f'图片尺寸为：{im.shape}，不是单通道，将只取第一个通道。')
#         im = im[:, :, 0] if np.array(im.shape).argmin() == 2 else im[0, :, :]
#     im = np.swapaxes(im, 0, 1)
#     # im_temp = im.copy()
#     im[im > 0] = 1
#
#     im = morphology.skeletonize(im)  # * im_temp
#     im = im.astype(np.uint8)
#
#     # extract a graph by placing vertices every THRESHOLD pixels, and at all intersections
#     vertices = []
#     edges = set()
#
#     def add_edge(src, dst):
#         if (src, dst) in edges or (dst, src) in edges:
#             return
#         elif src == dst:
#             return
#         edges.add((src, dst))
#
#     point_to_neighbors = {}
#     q = []
#     while True:
#         if len(q) > 0:
#             last_id, i, j = q.pop()
#             path = [vertices[last_id], (i, j)]
#             if im[i, j] == 0:
#                 continue
#             point_to_neighbors[(i, j)].remove(last_id)
#             if len(point_to_neighbors[(i, j)]) == 0:
#                 del point_to_neighbors[(i, j)]
#         else:
#             w = np.where(im > 0)
#             if len(w[0]) == 0:
#                 break
#             i, j = w[0][0], w[1][0]
#             last_id = len(vertices)
#             vertices.append((i, j))
#             path = [(i, j)]
#
#         while True:
#             im[i, j] = 0
#             neighbors = []
#             for oi in [-1, 0, 1]:
#                 for oj in [-1, 0, 1]:
#                     ni = i + oi
#                     nj = j + oj
#                     if 0 <= ni < im.shape[0] and 0 <= nj < im.shape[1] and im[ni, nj] > 0:
#                         neighbors.append((ni, nj))
#             if len(neighbors) == 1 and (i, j) not in point_to_neighbors:
#                 ni, nj = neighbors[0]
#                 path.append((ni, nj))
#                 i, j = ni, nj
#             else:
#                 if len(path) > 1:
#                     # path = rdp(path, 2)
#                     if len(path) > 2:
#                         for point in path[1:-1]:
#                             cur_id = len(vertices)
#                             vertices.append(point)
#                             add_edge(last_id, cur_id)
#                             last_id = cur_id
#                     neighbor_count = len(neighbors) + len(point_to_neighbors.get((i, j), []))
#                     if neighbor_count == 0 or neighbor_count >= 2:
#                         cur_id = len(vertices)
#                         vertices.append(path[-1])
#                         add_edge(last_id, cur_id)
#                         last_id = cur_id
#                 for ni, nj in neighbors:
#                     if (ni, nj) not in point_to_neighbors:
#                         point_to_neighbors[(ni, nj)] = set()
#                     point_to_neighbors[(ni, nj)].add(last_id)
#                     q.append((last_id, ni, nj))
#                 for neighbor_id in point_to_neighbors.get((i, j), []):
#                     add_edge(neighbor_id, last_id)
#                 break
#     neighbors = {}
#
#     vertex = vertices
#
#     for edge in edges:
#
#         nk1 = (vertex[edge[0]][1], vertex[edge[0]][0])
#         nk2 = (vertex[edge[1]][1], vertex[edge[1]][0])
#
#         if nk1 != nk2:
#             if nk1 in neighbors:
#                 if nk2 in neighbors[nk1]:
#                     pass
#                 else:
#                     neighbors[nk1].append(nk2)
#             else:
#                 neighbors[nk1] = [nk2]
#
#             if nk2 in neighbors:
#                 if nk1 in neighbors[nk2]:
#                     pass
#                 else:
#                     neighbors[nk2].append(nk1)
#             else:
#                 neighbors[nk2] = [nk1]
#
#     return neighbors


# def raster2line_shp(_mask: np.ndarray, geo_trans, project_ref, out_line_shp_path, class_infos):
#     if not isinstance(class_infos, dict):
#         class_infos = format_class_info(class_infos)
#     assert len(_mask.shape) == 2, '必须是单通道影像！'
#     shp_driver = ogr.GetDriverByName("ESRI Shapefile")
#     MyPath(out_line_shp_path).make_parent_dirs()
#     if MyPath(out_line_shp_path).exists():
#         shp_driver.DeleteDataSource(out_line_shp_path)
#     out_datasource = shp_driver.CreateDataSource(out_line_shp_path)
#     geo_trans = adjust_transform(geo_trans)
#     proj = osr.SpatialReference()
#     proj.ImportFromWkt(project_ref)
#     out_layer: ogr.Layer = out_datasource.CreateLayer('layer', srs=proj, geom_type=ogr.wkbMultiLineString)
#     create_field_defn(out_layer, class_fields := class_infos['class_fields'])
#     out_defn = out_layer.GetLayerDefn()  # 定义要素
#     out_feature_triangle = ogr.Feature(out_defn)
#     for i, class_info in enumerate(class_infos):
#         if i == 0:  # 背景
#             continue
#         cur_mask = (_mask == i).astype(np.uint8)
#         graph = generate_graph(cur_mask)
#         # 创建单个面
#         for n, v in graph.items():
#             for nei in v:
#                 line = ogr.Geometry(ogr.wkbLinearRing)  # 构建几何类型:线
#                 nx, ny = n[1], n[0]
#                 nx, ny = pixel2geo(nx, ny, geo_trans)
#                 line.AddPoint(nx, ny)  # 添加点01
#
#                 nei_x, nei_y = nei[1], nei[0]
#                 nei_x, nei_y = pixel2geo(nei_x, nei_y, geo_trans)
#                 line.AddPoint(nei_x, nei_y)  # 添加点02
#
#                 out_feature_triangle.SetGeometry(line)
#                 for field, field_type in class_fields.items():
#                     out_feature_triangle.SetField(field, class_info['class_names'].get(field, field_type(0)))
#                 out_layer.CreateFeature(out_feature_triangle)
#
#     out_datasource.Destroy()


def higher_bit2img8bit(img_higher_bit: (str, np.ndarray), img8bit_path=None, to_rgb=False, geo_transform=None,
                       proj=None, nodata=None):
    """
    遥感影像高位（16位，32位。。。）转8位影像
    Args:
        img_higher_bit: 待转换的图像
        img8bit_path: 转换后的8位图像保存地址
        to_rgb:
        geo_transform:
        proj:
        nodata:
    Returns:

    """
    if isinstance(img_higher_bit, str):  # 如果传进来的参数是图片路径，则先读取图片
        with OpenImage(img_higher_bit) as ds:
            nodata = ds.nodata  # 获取无效值
            geo_transform, proj = ds.geo_transform, ds.project  # 地理坐标变换系数和坐标系
            array_data = ds.img_array  # 读取array
    else:
        assert isinstance(img_higher_bit, np.ndarray)
        array_data = img_higher_bit
    if array_data.dtype == 'uint8':  # 如果本身就是8位，直接返回
        return array_data
    array_data[array_data == nodata] = 0  # 将无效值置为0
    if len(array_data.shape) == 2:  # 单通道灰度图，先升维，以便和彩色影像同样处理
        array_data = array_data[None, ...]
    counts, rows, cols = array_data.shape
    if to_rgb and counts > 3:  # to_rgb为真时，要求影像波段数大于3，只取前三个波段，RGB
        counts = 3
        array_data = array_data[:3, ...][::-1, ...]
    compress_data = np.zeros((counts, rows, cols), dtype=np.uint8)  # 初始化一个最终的8位影像

    for i in range(counts):  # 遍历各个波段，分别进行灰度值的压缩

        cur_band = array_data[i, ...]
        cut_min, cut_max = np.percentile(cur_band, 2), np.percentile(cur_band, 98)
        if cut_min == cut_max:
            continue
        compress_scale = (cut_max - cut_min) / 255  # 灰度值缩放比例
        cur_band[cur_band > cut_max] = cut_max  # 大于上区间的值直接过滤掉，下同
        cur_band[cur_band < cut_min] = cut_min
        compress_data[i, :, :] = (cur_band - cut_min) / compress_scale  # 将灰度值缩放到8位图的[0，255]之间

    if counts == 1:  # 单波段影像还原
        compress_data = compress_data[0, ...]

    if img8bit_path is not None:  # 设置了8位图保存路径则保存到文件，否则不保存
        write_image(compress_data, img8bit_path, geo_transform, proj)  # 将转换后的8位图写入文件
    return compress_data


# def batch_higher_bit2img8bit(img_higher_bit_dir, img8bit_dir, to_rgb=False, num_process=4):
#     """
#     批量遥感影像16位转8位影像
#     Args:
#         img_higher_bit_dir:
#         img8bit_dir:
#         to_rgb:
#         num_process:
#     Returns:
#
#     """
#
#     iterable_args = [(img_path.path, img_path.replace_parent(img8bit_dir).path, to_rgb) for img_path in
#                      MyPath(img_higher_bit_dir).iter_dir(ImageType.pattern())]
#     parallel_run_func(higher_bit2img8bit, iterable_args, num_workers=num_process, desc='批量转8位影像中。。。')


def resize_img(img: (str, Image.Image), target_size: tuple, save_path=None, pad=None):
    """
    使用Image对图片进行resize
    Args:
        img:
        target_size:(w,h)
        save_path:
        pad:
    Returns:

    """
    if isinstance(img, str):
        img: Image.Image = Image.open(img)
    assert isinstance(img, Image.Image)
    ori_w, ori_h = ori_size = img.size
    new_w, new_h = target_size
    if ori_size == target_size:
        return img
    if pad is None:  # 不填充，直接拉伸
        new_img = img.resize(target_size, Image.NEAREST)
    else:  # 保持横纵比，其余部分填充
        assert isinstance(pad, int), 'pad参数应该为整型数值，表示要填充的像素值'
        new_img = Image.new(img.mode, target_size, pad)
        scale_rate_w, scale_rate_h = new_w / ori_w, new_h / ori_h
        if scale_rate_w < scale_rate_h:
            img = img.resize((new_w, cur_h := int(ori_h * scale_rate_w)))
            new_img.paste(img, (0, (new_h - cur_h) // 2))
        else:
            img = img.resize((cur_w := int(ori_w * scale_rate_h), new_h))
            new_img.paste(img, ((new_w - cur_w) // 2, 0))
    if save_path is not None:
        MyPath(save_path).make_parent_dirs()
        new_img.save(save_path)
        print(f'{save_path}尺寸已从{ori_size}resize为{target_size}')
    return new_img


def resize_img_group(images: (list, tuple), target_size: tuple, save_paths=None, pad=None):
    """
    使用Image对一组图片进行同步resize
    Args:
        images:
        target_size:
        save_paths:
        pad:
    Returns:

    """
    if isinstance(images, str):  # 单张图片
        images = [images]
    assert isinstance(images, (list, tuple))
    results = []
    if save_paths is not None:
        if isinstance(save_paths, str):
            save_paths = [save_paths]
        assert isinstance(save_paths, (list, tuple)) and len(images) == len(save_paths)
    else:
        save_paths = [None] * len(images)

    for img, save_path in zip(images, save_paths):
        img = resize_img(img, target_size, save_path, pad)
        results.append(img)

    return results


# def batch_resize_img(img_dirs: (str, list, tuple), target_size=(512, 512), num_workers=4, pad=None):
#     """
#     批量对目录下得所有图片进行resize到统一尺寸
#     Args:
#         img_dirs: 图片目录列表
#         target_size: 目标尺寸
#         num_workers: 多进程个数
#         pad:
#     Returns:
#
#     """
#     if isinstance(img_dirs, str):
#         img_dirs = [img_dirs]
#     list_suffix = get_list_suffix_from_dirs(img_dirs)
#     assert len(list_suffix) == len(img_dirs)
#     iterable_args = []
#     for img_path in MyPath(img_dirs[0]).iter_dir(ImageType.pattern()):
#         img_path.real_suffix = list_suffix[0]
#         images = [img_path.path]
#         for img_dir, suffix in list(zip(img_dirs, list_suffix))[1:]:
#             img_path_content = img_path.replace_parent(img_dir).replace_suffix(suffix)
#             images.append(img_path_content.path)
#         iterable_args.append((images, target_size, images, pad))
#     parallel_run_func(resize_img_group, iterable_args, num_workers=num_workers, desc='批量resize图片中。。。')


# gdal版本
# def extract_by_shp_gdal(in_shp_path, in_raster_path, out_raster_path='mem', nodata=0):
#     """根据矢量的不规则范围对栅格影像进行裁剪，非范围内的赋无效值"""
#     input_raster = gdal.Open(in_raster_path)
#
#     # 利用gdal.Warp进行裁剪
#     # https://gdal.org/api/python/osgeo.gdal.html#osgeo.gdal.Warp
#     MyPath(out_raster_path).make_parent_dirs()
#     result: gdal.Dataset = gdal.Warp(
#         out_raster_path,
#         input_raster,
#         format='GTiff',
#         cutlineDSName=in_shp_path,  # 用于裁剪的矢量
#         cropToCutline=True,  # 是否使用cut_lineDSName的extent作为输出的界线
#         dstNodata=nodata  # 输出数据的nodata值
#     )
#     result.FlushCache()
#     out_image_array = result.ReadAsArray()
#     out_transform = result.GetGeoTransform()
#     del result
#     return out_image_array, out_transform

# rasterio版本


def extract_by_shp(in_shp, in_raster_path, out_raster_path=None):
    """
    根据矢量的不规则范围对栅格影像进行裁剪，非范围内的赋无效值
    :param in_shp: shp文件路径或者wkt格式的polygon
    :param in_raster_path: 图片路径
    :param out_raster_path: 剪裁图片保存路径
    :return:
    """
    assert isinstance(in_shp, str), (error_in_shp := 'inshp参数只能是shp文件路径或者wkt格式的polygon字符串！')
    with OpenImage(in_raster_path, backend='rasterio') as src:
        if MyPath(in_shp).is_file():
            # x_min, y_min, x_max, y_max = src.img_data.bounds
            # wkt = f"POLYGON (({x_min} {y_min},{x_max} {y_min},{x_max} {y_max},{x_min} {y_max}))"
            # spatial_filter = ogr.CreateGeometryFromWkt(wkt)
            with read_shp(in_shp, spatial_filter=None) as layer:
                in_shp = [json.loads(feature.geometry().ExportToJson()) for feature in layer]
        elif in_shp.startswith('POLYGON'):
            geometry_bound: ogr.Geometry = ogr.CreateGeometryFromWkt(in_shp)
            in_shp = [json.loads(geometry_bound.ExportToJson())]
        else:
            raise ValueError(error_in_shp)
        nodata = src.nodata or 0
        out_image, out_transform = mask(src.img_data, in_shp, nodata=nodata, crop=True)
        geo_project = src.project

    # 保存
    if out_raster_path is not None:
        MyPath(out_raster_path).make_parent_dirs()
        write_image(out_image, out_raster_path, out_transform, geo_project, nodata=nodata)
    return out_image, out_transform


def rgb_label2mask(rgb_label_path: str, rgb_mapping: dict, out_path=None, palette=None):
    """
    rgb真彩色标签转单通道训练标签
    :param rgb_label_path: rgb标签图片路径，
    :param rgb_mapping: rgb颜色映射字典，{(r,g,b):class_id}，key,value分别对应rgb值和类别id值
    :param out_path: 标签保存路径，默认为空
    :param palette: 伪彩rgb映射列表
    :return:
    """
    rgb_label = read_image_as_array(rgb_label_path, backend='pil')
    assert len(rgb_label.shape) == 3, '必须是真彩色标签图像！'
    out_label = np.zeros(rgb_label.shape[:2], dtype=np.uint8)
    for rgb, class_id in rgb_mapping.items():
        r, g, b = rgb[:3]
        mask_rgb = (rgb_label[..., 0] == r) & (rgb_label[..., 1] == g) & (rgb_label[..., 2] == b)
        out_label[mask_rgb] = class_id
    if out_path is None:
        return out_label
    MyPath(out_path).make_parent_dirs()
    with OpenImage(out_path, 'w', backend='pil') as out:
        out.set_img_array(out_label)
        if palette is None:
            palette = {class_id: rgb for rgb, class_id in rgb_mapping.items()}
            palette = np.array(list(dict(sorted(palette.items())).values())).flatten().tolist()
        out.set_color_map(palette)
    return out_label


def batch_rgb_label2mask(rgb_label_dir, rgb_mapping: dict, out_dir, palette=None):
    for label_path in tqdm(MyPath(rgb_label_dir).iter_dir()):
        out_path = label_path.replace_parent(out_dir)
        rgb_label2mask(str(label_path), rgb_mapping, str(out_path), palette)


# def seg_result2change_result(seg_result1, seg_result2, dict_replace_values=None):
#     if dict_replace_values is None:
#         dict_replace_values = {255: 2}
#     if isinstance(seg_result1, str):
#         seg_result1 = read_image_as_array(seg_result1)
#     if isinstance(seg_result2, str):
#         seg_result2 = read_image_as_array(seg_result2)
#     assert isinstance(seg_result1, np.ndarray) and isinstance(seg_result2, np.ndarray)
#     if seg_result1.shape != seg_result2.shape:
#         seg_result2 = mmcv.imresize(seg_result2, seg_result1.shape[::-1])
#
#     change_result = seg_result2 - seg_result1
#     change_result = sieve(change_result, size=100)
#     change_result = img_replace_values(change_result, dict_replace_values)
#     return change_result


def img_replace_values(img: (np.ndarray, str), dict_replace_values: dict, save_path=None) -> np.ndarray:
    """
    替换img_array中的值
    Args:
        img: 图片路径或者array
        dict_replace_values: 替换字典，格式例子：{255:0},{2:(1,0)},{(2,3):(1,0)},{(255,1):0,(2,3):1,(4,5):2}
        save_path: 保存路径，默认为None，不保存。

    Returns:

    """

    assert isinstance(dict_replace_values,
                      dict), 'replace_values参数类型应该为字典类型！key为原始值(多个值可以用元组)，value为新值(' \
                             '长度为2的元组，第一个元素为key值对应的新值，第二个值表示除key值外的其他值对应的新值,默认为None表示不变)'
    if isinstance(img, str):
        img = read_image_as_array(img)
    assert isinstance(img, np.ndarray), f'img参数类型错误！'
    for ori_values, new_values in dict_replace_values.items():
        if isinstance(ori_values, int):  # key为单个值整数时，先变为元组
            ori_values = (ori_values,)
        assert isinstance(ori_values, tuple)
        if isinstance(new_values, int):
            new_values = (new_values, None)
        assert isinstance(new_values, (tuple, list)) and len(new_values) >= 2
        if len(ori_values) == 1:  # 要改的值只有一个时，不适用于reduce函数
            arg_mask = img == ori_values[0]
        else:  # 同时修改多个值，用reduce函数生成掩码
            arg_mask = reduce(lambda x, y: (img == x) | (img == y), ori_values)
        img = np.where(arg_mask, new_values[0],
                       new_values[1] if new_values[1] is not None else img)  # 字典的value元组第二值为None时，表示其他值不变
    if save_path is not None:
        write_image(img.astype(np.uint8), save_path)
    return img


def mask_valid(img_array):
    """
    计算有效值掩码,即多个通道不全为0或者不全为255的为有效
    :param img_array:
    :return:
    """
    mask_nodata1 = img_array != 0
    mask_nodata2 = img_array != 255
    if len(img_array.shape) == 3:  # 如果是rgb影像，需要对多个通道进行判断有效值，并转为单通道掩码，与最终预测结果shape保持一致
        mask_nodata1 = reduce(lambda _x, _y: _x | _y, mask_nodata1)  # 多个通道取并，即只要任意一个通道不是无效值，就判断为有效
        mask_nodata2 = reduce(lambda _x, _y: _x | _y, mask_nodata2)
    return mask_nodata1 * mask_nodata2
